package org.freespace.service.impl;

import cn.hutool.core.io.FileUtil;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.baomidou.mybatisplus.core.toolkit.Wrappers;
import com.blazemeter.jmeter.threads.arrivals.ArrivalsThreadGroup;
import com.blazemeter.jmeter.threads.concurrency.ConcurrencyThreadGroup;
import org.apache.jmeter.control.LoopController;
import org.apache.jmeter.save.SaveService;
import org.apache.jmeter.threads.ThreadGroup;
import org.apache.jorphan.collections.HashTree;
import org.freespace.model.JmxSample;
import org.freespace.model.JmxScript;
import org.freespace.model.JmxThreadgroup;
import org.freespace.mapper.JmxThreadgroupMapper;
import org.freespace.service.IJmxSampleService;
import org.freespace.service.IJmxScriptService;
import org.freespace.service.IJmxThreadgroupService;
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import org.freespace.vo.JmeterConcurrencyThreadVO;
import org.freespace.vo.JmeterThreadVO;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.time.LocalDateTime;
import java.util.List;
import java.util.Map;
import java.util.Set;

/**
 * <p>
 * 压测机器表 服务实现类
 * </p>
 *
 * @author liuqi
 * @since 2023-05-12
 */
@Service
public class JmxThreadgroupServiceImpl extends ServiceImpl<JmxThreadgroupMapper, JmxThreadgroup> implements IJmxThreadgroupService {

    @Autowired
    IJmxSampleService sampleService;
    @Autowired
    IJmxScriptService scriptService;

    @Autowired
    IJmxThreadgroupService threadgroupService;
    @Override
    public void deleteJmxThreadgroupByJmxId(Integer jmxId){
        QueryWrapper<JmxThreadgroup> wrapper = new QueryWrapper<>();
        wrapper.eq("jmx_id", jmxId);
        remove(wrapper);
    }

    @Override
    public List<JmxThreadgroup> list(Integer jmxId) {
        LambdaQueryWrapper<JmxThreadgroup> w = Wrappers.<JmxThreadgroup>lambdaQuery().eq(JmxThreadgroup::getJmxId, jmxId);
        return list(w);
    }

    @Override
    public List<JmxSample> samplelist(Integer jmxId, Integer threadGroupId) {
        LambdaQueryWrapper<JmxSample> w = Wrappers.<JmxSample>lambdaQuery().eq(JmxSample::getJmxId, jmxId).eq(JmxSample::getThreadgroupId, threadGroupId);
        return sampleService.list(w);
    }

    @Override
    public void enableThreadGroup(Integer threadGroupId, Integer isEnable) throws Exception {
        Boolean b = false;
        if(isEnable == 1){
            b = true;
        }

        JmxThreadgroup threadgroup = getById(threadGroupId);
        Map<String, Object> attrMap = JSONUtil.parseObj(threadgroup.getAttribute());
        if (threadgroup.getType().equals("ThreadGroup")) {
            enableJmxThread(threadgroup, attrMap, b);
        } else if (threadgroup.getType().equals("ConcurrencyThreadGroup")) {
            enableConcurrencyThread(threadgroup, attrMap, b);
        }
        threadgroup.setStatus(isEnable);
        updateById(threadgroup);
    }

    private void enableJmxThread(JmxThreadgroup threadgroup, Map<String, Object> attrMap, Boolean b) throws Exception {
        JmxScript jmxScript = scriptService.getById(threadgroup.getJmxId());
        String jmxPath = jmxScript.getPath();
        File localFile = scriptService.readFSScriptFileToLocal(jmxScript);

        scriptService.initProperties();
        HashTree jmxTree = SaveService.loadTree(localFile);
        Set<Object> keys = jmxTree.keySet();

        for (Object key : keys) {
            if (key instanceof ThreadGroup) {
                String testname = ((ThreadGroup) key).getPropertyAsString(ThreadGroup.NAME);
                if (testname.equals(threadgroup.getTestName())) {
                    JSONObject attMap = JSONUtil.parseObj(threadgroup.getAttribute());
                    Long numThreads = ((ThreadGroup) key).getPropertyAsLong(ThreadGroup.NUM_THREADS);
                    Long duration = ((ThreadGroup) key).getPropertyAsLong(ThreadGroup.DURATION);
                    Long rampTime = ((ThreadGroup) key).getPropertyAsLong(ThreadGroup.RAMP_TIME);
                    Boolean sch = ((ThreadGroup) key).getPropertyAsBoolean(ThreadGroup.SCHEDULER);
                    if (sch.equals(attMap.get("scheduler"))
                        && rampTime.intValue() == (Integer) attMap.get("rampup")
                            && duration.intValue() == (Integer) attMap.get("duration")
                            && numThreads.intValue() == (Integer) attMap.get("threads")) {
                        ((ThreadGroup) key).setEnabled(b);
                        SaveService.saveTree(jmxTree, new FileOutputStream(localFile));

                        scriptService.overrideRemoteFile(threadgroup.getJmxId(), localFile.getName(), FileUtil.getInputStream(localFile));
                    }
                }
                break;
            }
        }

        for (HashTree item : jmxTree.values()) {
            Set treeKeys = item.keySet();
            for (Object key : treeKeys) {
                if (key instanceof ThreadGroup) {
                    String testname = ((ThreadGroup) key).getPropertyAsString(ThreadGroup.NAME);
                    if (testname.equals(threadgroup.getTestName())) {
                        JSONObject attMap = JSONUtil.parseObj(threadgroup.getAttribute());
                        Long numThreads = ((ThreadGroup) key).getPropertyAsLong(ThreadGroup.NUM_THREADS);
                        Long duration = ((ThreadGroup) key).getPropertyAsLong(ThreadGroup.DURATION);
                        Long rampTime = ((ThreadGroup) key).getPropertyAsLong(ThreadGroup.RAMP_TIME);
                        Boolean sch = ((ThreadGroup) key).getPropertyAsBoolean(ThreadGroup.SCHEDULER);
                        if (sch.equals(attMap.get("scheduler"))
                                && rampTime.intValue() == (Integer) attMap.get("rampup")
                                && duration.intValue() == (Integer) attMap.get("duration")
                                && numThreads.intValue() == (Integer) attMap.get("threads")) {
                            ((ThreadGroup) key).setEnabled(b);
                            SaveService.saveTree(jmxTree, new FileOutputStream(localFile));

                            scriptService.overrideRemoteFile(threadgroup.getJmxId(), localFile.getName(), FileUtil.getInputStream(localFile));
                        }
                    }
                    break;
                }
            }
        }
    }

    private void enableConcurrencyThread(JmxThreadgroup threadgroup, Map<String, Object> attrMap, Boolean b) throws Exception{
        JmxScript jmxScript = scriptService.getById(threadgroup.getJmxId());
        String jmxPath = jmxScript.getPath();
        File localFile = scriptService.readFSScriptFileToLocal(jmxScript);

        scriptService.initProperties();
        HashTree jmxTree = SaveService.loadTree(localFile);
        Set<Object> keys = jmxTree.keySet();

        for (Object key : keys) {
            if (key instanceof ConcurrencyThreadGroup) {
                String testname = ((ConcurrencyThreadGroup) key).getPropertyAsString(ThreadGroup.NAME);
                if (testname.equals(threadgroup.getTestName())) {
                    JSONObject attMap = JSONUtil.parseObj(threadgroup.getAttribute());
                    Long targetLevel = ((ConcurrencyThreadGroup) key).getPropertyAsLong(ConcurrencyThreadGroup.TARGET_LEVEL);
                    Long hold = ((ConcurrencyThreadGroup) key).getPropertyAsLong(ConcurrencyThreadGroup.HOLD);
                    Long rampUp = ((ConcurrencyThreadGroup) key).getPropertyAsLong(ConcurrencyThreadGroup.RAMP_UP);
                    Long steps = ((ConcurrencyThreadGroup) key).getPropertyAsLong(ConcurrencyThreadGroup.STEPS);
                    if (targetLevel.intValue() == (Integer) attMap.get("targetConcurrency")
                            && rampUp.intValue() == (Integer) attMap.get("rampUpTime")
                            && hold.intValue() == (Integer) attMap.get("holdTime")
                            && steps.intValue() == (Integer) attMap.get("steps") ) {
                        ((ThreadGroup) key).setEnabled(b);
                        SaveService.saveTree(jmxTree, new FileOutputStream(localFile));

                        scriptService.overrideRemoteFile(threadgroup.getJmxId(), localFile.getName(), FileUtil.getInputStream(localFile));
                    }
                }
            }
        }

        for (HashTree item : jmxTree.values()) {
            Set treeKeys = item.keySet();
            for (Object key : treeKeys) {
                if (key instanceof ConcurrencyThreadGroup) {
                    String testname = ((ConcurrencyThreadGroup) key).getPropertyAsString(ThreadGroup.NAME);
                    if (testname.equals(threadgroup.getTestName())) {
                        JSONObject attMap = JSONUtil.parseObj(threadgroup.getAttribute());
                        Long targetLevel = ((ConcurrencyThreadGroup) key).getPropertyAsLong(ConcurrencyThreadGroup.TARGET_LEVEL);
                        Long hold = ((ConcurrencyThreadGroup) key).getPropertyAsLong(ConcurrencyThreadGroup.HOLD);
                        Long rampUp = ((ConcurrencyThreadGroup) key).getPropertyAsLong(ConcurrencyThreadGroup.RAMP_UP);
                        Long steps = ((ConcurrencyThreadGroup) key).getPropertyAsLong(ConcurrencyThreadGroup.STEPS);
                        if (targetLevel.intValue() == (Integer) attMap.get("targetConcurrency")
                                && rampUp.intValue() == (Integer) attMap.get("rampUpTime")
                                && hold.intValue() == (Integer) attMap.get("holdTime")
                                && steps.intValue() == (Integer) attMap.get("steps") ) {
                            ((ThreadGroup) key).setEnabled(b);
                            SaveService.saveTree(jmxTree, new FileOutputStream(localFile));

                            scriptService.overrideRemoteFile(threadgroup.getJmxId(), localFile.getName(), FileUtil.getInputStream(localFile));
                        }
                    }
                }
            }
        }
    }

    @Override
    public void updateJmxThread(JmeterThreadVO req) throws Exception {
        JmxScript jmxScript = scriptService.getById(req.getJmxId());
        String jmxPath = jmxScript.getPath();
        File localFile = scriptService.readFSScriptFileToLocal(jmxScript);

        scriptService.initProperties();
        HashTree jmxTree = SaveService.loadTree(localFile);
        Set<Object> keys = jmxTree.keySet();

        JmxThreadgroup threadgroup = threadgroupService.getById(req.getId());

        for (Object key : keys) {
            if (key instanceof ThreadGroup) {
                String testname = ((ThreadGroup) key).getPropertyAsString(ThreadGroup.NAME);
                if (testname.equals(threadgroup.getTestName())) {
                    if (((ThreadGroup) key).isEnabled()) {
                        updateThreadGroup(req, localFile, jmxTree, threadgroup, (ThreadGroup) key);
                    }
                    break;
                }
            }
        }

        for (HashTree item : jmxTree.values()) {
            Set treeKeys = item.keySet();
            for (Object key : treeKeys) {
                if (key instanceof ThreadGroup) {
                    String testname = ((ThreadGroup) key).getPropertyAsString(ThreadGroup.NAME);
                    if (testname.equals(threadgroup.getTestName())) {
                        if (((ThreadGroup) key).isEnabled()) {
                            updateThreadGroup(req, localFile, jmxTree, threadgroup, (ThreadGroup) key);
                        }
                    }
                    break;
                }
            }
        }
    }

    private void updateThreadGroup(JmeterThreadVO req, File localFile, HashTree jmxTree, JmxThreadgroup threadgroup, ThreadGroup key) throws IOException {
        JSONObject attMap = JSONUtil.parseObj(threadgroup.getAttribute());
        if (req.getFlag() == 1) {
            if (req.getUserCount() > 0) {
                key.setProperty(ThreadGroup.NUM_THREADS, req.getUserCount());
                attMap.put("threads", req.getUserCount());
            }
            if (req.getDurationTime() > 0) {
                key.setProperty(ThreadGroup.DURATION, req.getDurationTime());
                attMap.put("duration", req.getDurationTime());
            }
            if (req.getRampTime() > 0) {
                key.setProperty(ThreadGroup.RAMP_TIME, req.getRampTime());
                attMap.put("rampup", req.getRampTime());
            }
            key.setProperty(ThreadGroup.SCHEDULER, true);
            attMap.put("scheduler", true);

            LoopController controller = (LoopController) key.getSamplerController();
            if (controller != null) {
                controller.setProperty(LoopController.LOOPS, -1);
                attMap.put("loops", -1);
                controller.setContinueForever(true);
                controller.setProperty("LoopController.continue_forever", true);
                attMap.put("isForever", true);
                key.setSamplerController(controller);
            }
            SaveService.saveTree(jmxTree, new FileOutputStream(localFile));
            scriptService.overrideRemoteFile(threadgroup.getJmxId(), localFile.getName(), FileUtil.getInputStream(localFile));
        } else {
            if (req.getUserCount() > 0) {
                key.setProperty(ThreadGroup.NUM_THREADS, req.getUserCount());
                attMap.put("threads", req.getUserCount());
            }
            if (req.getRampTime() > 0) {
                key.setProperty(ThreadGroup.RAMP_TIME, req.getRampTime());
                attMap.put("rampup", req.getRampTime());
            }
            key.setProperty(ThreadGroup.SCHEDULER, false);
            attMap.put("scheduler", false);

            LoopController controller = (LoopController) key.getSamplerController();
            if (controller != null) {
                controller.setProperty(LoopController.LOOPS, req.getLoopCount());
                attMap.put("loops", req.getLoopCount());
                controller.setContinueForever(true);
                controller.setProperty("LoopController.continue_forever", true);
                attMap.put("isForever", true);
                key.setSamplerController(controller);
            }
            SaveService.saveTree(jmxTree, new FileOutputStream(localFile));
            scriptService.overrideRemoteFile(threadgroup.getJmxId(), localFile.getName(), FileUtil.getInputStream(localFile));
        }
        threadgroup.setLastUpdateTime(LocalDateTime.now());
        threadgroup.setAttribute(JSONUtil.toJsonStr(attMap));
        threadgroupService.updateById(threadgroup);
    }

    private void updateConcurrencyThreadGroup(JmeterConcurrencyThreadVO req, File localFile, HashTree jmxTree, JmxThreadgroup threadgroup, ConcurrencyThreadGroup key) throws IOException {
        JSONObject attMap = JSONUtil.parseObj(threadgroup.getAttribute());
        if (req.getTargetConcurrency() != null) {
            ((ConcurrencyThreadGroup)key).setProperty(ConcurrencyThreadGroup.TARGET_LEVEL, req.getTargetConcurrency());
            attMap.put("targetConcurrency", req.getTargetConcurrency());
        }

        if (req.getHoldTime() != null) {
            ((ConcurrencyThreadGroup)key).setProperty(ConcurrencyThreadGroup.HOLD, req.getTargetConcurrency());
            attMap.put("holdTime", req.getHoldTime());
        }

        if (req.getRampUpTime() != null) {
            ((ConcurrencyThreadGroup)key).setProperty(ConcurrencyThreadGroup.RAMP_UP, req.getRampUpTime());
            attMap.put("rampUpTime", req.getRampUpTime());
        }

        if (req.getSteps() != null) {
            ((ConcurrencyThreadGroup)key).setProperty(ConcurrencyThreadGroup.STEPS, req.getSteps());
            attMap.put("steps", req.getSteps());
        }

        if (req.getUnit() != null) {
            ((ConcurrencyThreadGroup)key).setProperty(ConcurrencyThreadGroup.UNIT, req.getUnit());
            attMap.put("Unit", req.getUnit());
        }

        SaveService.saveTree(jmxTree, new FileOutputStream(localFile));
        scriptService.overrideRemoteFile(threadgroup.getJmxId(), localFile.getName(), FileUtil.getInputStream(localFile));

        threadgroup.setLastUpdateTime(LocalDateTime.now());
        threadgroup.setAttribute(JSONUtil.toJsonStr(attMap));
        threadgroup.setUnit(req.getUnit());
        threadgroupService.updateById(threadgroup);
    }

    private void updateArrivalsThreadGroup(JmeterConcurrencyThreadVO req, File localFile, HashTree jmxTree, JmxThreadgroup threadgroup, ArrivalsThreadGroup key) throws IOException {
        JSONObject attMap = JSONUtil.parseObj(threadgroup.getAttribute());
        if (req.getTargetConcurrency() != null) {
            key.setProperty(ConcurrencyThreadGroup.TARGET_LEVEL, req.getTargetConcurrency());
            attMap.put("targetConcurrency", req.getTargetConcurrency());
        }

        if (req.getHoldTime() != null) {
            key.setProperty(ConcurrencyThreadGroup.HOLD, req.getTargetConcurrency());
            attMap.put("holdTime", req.getHoldTime());
        }

        if (req.getRampUpTime() != null) {
            key.setProperty(ConcurrencyThreadGroup.RAMP_UP, req.getRampUpTime());
            attMap.put("rampUpTime", req.getRampUpTime());
        }

        if (req.getSteps() != null) {
            key.setProperty(ConcurrencyThreadGroup.STEPS, req.getSteps());
            attMap.put("steps", req.getSteps());
        }

        if (req.getUnit() != null) {
            key.setProperty(ConcurrencyThreadGroup.UNIT, req.getUnit());
            attMap.put("Unit", req.getUnit());
        }

        SaveService.saveTree(jmxTree, new FileOutputStream(localFile));
        scriptService.overrideRemoteFile(threadgroup.getJmxId(), localFile.getName(), FileUtil.getInputStream(localFile));

        threadgroup.setLastUpdateTime(LocalDateTime.now());
        threadgroup.setAttribute(JSONUtil.toJsonStr(attMap));
        threadgroup.setUnit(req.getUnit());
        threadgroupService.updateById(threadgroup);
    }

    @Override
    public void updateJmxArrivalsThread(JmeterConcurrencyThreadVO req) throws Exception{
        JmxScript jmxScript = scriptService.getById(req.getJmxId());
        String jmxPath = jmxScript.getPath();
        File localFile = scriptService.readFSScriptFileToLocal(jmxScript);

        scriptService.initProperties();
        HashTree jmxTree = SaveService.loadTree(localFile);
        Set<Object> keys = jmxTree.keySet();

        JmxThreadgroup threadgroup = threadgroupService.getById(req.getId());

        for (Object key : keys) {
            if (key instanceof ArrivalsThreadGroup) {
                String testname = ((ArrivalsThreadGroup) key).getPropertyAsString(ThreadGroup.NAME);
                if (testname.equals(threadgroup.getTestName())) {
                    if (((ArrivalsThreadGroup) key).isEnabled()) {
                        updateArrivalsThreadGroup(req, localFile, jmxTree, threadgroup, (ArrivalsThreadGroup) key);
                        break;
                    }
                }
            }
        }

        for (HashTree item : jmxTree.values()) {
            Set treeKeys = item.keySet();
            for (Object key : treeKeys) {
                if (key instanceof ArrivalsThreadGroup) {
                    String testname = ((ArrivalsThreadGroup) key).getPropertyAsString(ThreadGroup.NAME);
                    if (testname.equals(threadgroup.getTestName())) {
                        if (((ArrivalsThreadGroup) key).isEnabled()) {
                            updateArrivalsThreadGroup(req, localFile, jmxTree, threadgroup, (ArrivalsThreadGroup) key);
                            break;
                        }
                    }
                }
            }
        }
    }

    @Override
    public void updateJmxConcurrencyThread(JmeterConcurrencyThreadVO req) throws Exception {
        JmxScript jmxScript = scriptService.getById(req.getJmxId());
        String jmxPath = jmxScript.getPath();
        File localFile = scriptService.readFSScriptFileToLocal(jmxScript);

        scriptService.initProperties();
        HashTree jmxTree = SaveService.loadTree(localFile);
        Set<Object> keys = jmxTree.keySet();

        JmxThreadgroup threadgroup = threadgroupService.getById(req.getId());

        for (Object key : keys) {
            if (key instanceof ConcurrencyThreadGroup) {
                String testname = ((ConcurrencyThreadGroup) key).getPropertyAsString(ThreadGroup.NAME);
                if (testname.equals(threadgroup.getTestName())) {
                    if (((ConcurrencyThreadGroup) key).isEnabled()) {
                        updateConcurrencyThreadGroup(req, localFile, jmxTree, threadgroup, (ConcurrencyThreadGroup) key);
                    }
                    break;
                }
            }
        }

        for (HashTree item : jmxTree.values()) {
            Set treeKeys = item.keySet();
            for (Object key : treeKeys) {
                if (key instanceof ConcurrencyThreadGroup) {
                    String testname = ((ConcurrencyThreadGroup) key).getPropertyAsString(ThreadGroup.NAME);
                    if (testname.equals(threadgroup.getTestName())) {
                        if (((ConcurrencyThreadGroup) key).isEnabled()) {
                            updateConcurrencyThreadGroup(req, localFile, jmxTree, threadgroup, (ConcurrencyThreadGroup) key);
                        }
                    }
                    break;
                }
            }
        }
    }
}
